#include <dlfcn.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include "memcheck.h"

static int (*real_start_main)(int (*main_fn)(int, char **, char **), int, char **, void(*)(void), void(*)(void), void(*)(void), void *);
static int (*real_start_main_musl)(int (*main_fn)(int, char **, char **), int, char **, void(*)(void), void(*)(void), void(*)(void));

static int (*g_real_main)(int, char **, char **);

static int memcheck_main_wrapper(int argc, char **argv, char **envp)
{
    if (dbi_init() != 0) {
        return g_real_main(argc, argv, envp);
    }
    void *tramp = dbi_translate_function((void *)g_real_main);
    if (!tramp) {
        return g_real_main(argc, argv, envp);
    }
    int (*tramp_main)(int, char **, char **) = (int (*)(int, char **, char **))tramp;
    return tramp_main(argc, argv, envp);
}

int __libc_start_main(int (*main_fn)(int, char **, char **), int argc, char **ubp_av,
                      void (*init)(void), void (*fini)(void), void (*rtld_fini)(void), void *stack_end)
{
    if (!real_start_main) {
        real_start_main = (int (*)(int (*)(int,char **,char **), int, char **, void(*)(void), void(*)(void), void(*)(void), void *))
            dlsym(RTLD_NEXT, "__libc_start_main");
    }
    g_real_main = main_fn;
    return real_start_main(memcheck_main_wrapper, argc, ubp_av, init, fini, rtld_fini, stack_end);
}

// __attribute__((weak))
// int __libc_start_main_stage2(int (*main_fn)(int, char **, char **), int argc, char **argv,
//                              void (*init)(void), void (*fini)(void), void (*ldso_fini)(void))
// {
//     if (!real_start_main_musl) {
//         real_start_main_musl = (int (*)(int (*)(int,char **,char **), int, char **, void(*)(void), void(*)(void), void(*)(void)))
//             dlsym(RTLD_NEXT, "__libc_start_main_stage2");
//     }
//     if (!real_start_main_musl) {
//         return main_fn(argc, argv, NULL);
//     }
//     g_real_main = main_fn;
//     return real_start_main_musl(memcheck_main_wrapper, argc, argv, init, fini, ldso_fini);
// }
